
/* Copyright (c) Gong Chengming, NARI-TECH, 2011 */

#ifndef _LECTURE4_MATRIX_H
#define _LECTURE4_MATRIX_H

#include <cstddef>
#include <iostream>
#include <vector>

/*
  Column major dense matrix class template.
*/

template<class T>
class Matrix {
 public:
  Matrix() {
    Resize(0, 0);
  }
  Matrix(size_t rows, size_t cols) {
    Resize(rows, cols);
  }
  ~Matrix() {}

  Matrix<T>& Resize(size_t rows, size_t cols) {
    _rows = rows;
    _cols = cols;
    _data.resize(_rows * _cols);
    return *this;
  }

  size_t Rows() const {
    return _rows;
  }

  size_t Cols() const {
    return _cols;
  }

  T& operator() (size_t i, size_t j) {
//    std::cout<<"i = "<<i<<", j = "<<j<<", sub = "<<Rows() * j + i<<std::endl;
    return _data[Rows() * j + i];
  }

  const T& operator() (size_t i, size_t j) const {
    return _data[Rows() * j + i];
  }

  void Assign(T* v_vec) {
    _data.assign(v_vec, v_vec + Rows() * Cols());
  }

 private:
  size_t _rows;
  size_t _cols;
  std::vector<T> _data;
};

template<class T>
std::ostream & operator << (std::ostream& os, Matrix<T>& m) {
  for (size_t r = 0; r < m.Rows(); ++r) {
    for (size_t c = 0; c < m.Cols(); ++c) {
      os << m(r,c) << '\t';
    }
    os << "\n";
  }
  return os;
}

/*
 A squre matrix class template.
*/

template<class T>
class Square : private Matrix<T> {
 public:
  Square() {
    SetDim(0);
  }
  Square(size_t dim) {
    SetDim(dim);
  }
  ~Square() {}

  Square<T>& SetDim(size_t dim) {
    Matrix<T>::Resize(dim, dim);
    return *this;
  }

  size_t Dim() const {
    return Matrix<T>::Cols();
  }

  T& operator()(size_t i, size_t j) {
    return Matrix<T>::operator()(i, j);
  }

  const T& operator() (size_t i, size_t j) const {
    return Matrix<T>::operator()(i, j);
  }

  void Assign(T* v_vec) {
    Matrix<T>::Assign(v_vec);
  }
};

template<class T>
std::ostream & operator << (std::ostream& os, Square<T>& m) {
  for (size_t r = 0; r < m.Dim(); ++r) {
    for (size_t c = 0; c < m.Dim(); ++c) {
//      std::cout<<"r = "<<r<<", c = "<<c<<std::endl;
      os << m(r,c) << '\t';
    }
    os << "\n";
  }
  return os;
}

#endif  // _LECTURE4_MATRIX_H
